import openpyxl
import xlsxwriter
from PIL import Image, ImageDraw, ImageFont
import os

def create_watermark_image(text, output_path, image_size=(600, 180)):
    """
    创建水印图像，并保存到文件。
    """
    width, height = image_size  # 水印图像的大小
    img = Image.new('RGBA', (width, height), (255, 255, 255, 0))  # 透明背景
    # 获取字体
    try:
        font = ImageFont.truetype("simkai.ttf", 72)
    except IOError:
        print("默认字体")
        font = ImageFont.load_default()  # 如果没有arial字体，则使用默认字体
    draw = ImageDraw.Draw(img)
    # 使用textbbox来获取文本的边界框
    bbox = draw.textbbox((0, 0), text, font=font)
    text_width = bbox[2] - bbox[0]  # 计算文本的宽度
    text_height = bbox[3] - bbox[1]  # 计算文本的高度
    # 将水印文本居中绘制，文本颜色改为灰色（255, 255, 255, 128为半透明灰色）
    draw.text(((width - text_width) / 2, (height - text_height) / 2), text, fill=(169, 169, 169, 128), font=font)  # 半透明灰色
    # 旋转图像60度
    img = img.rotate(60, expand=True)
    # 调整水印图像大小
    img.thumbnail((image_size[0], image_size[1]), Image.Resampling.LANCZOS)  # 使用 LANCZOS 代替 ANTIALIAS
    # 保存图像
    img.save(output_path)



def add_watermark_to_excel(input_file, output_file, watermark_text):
    # 获取文件夹路径
    folder_path = os.path.dirname(output_file)
    # 检查文件夹是否存在，如果不存在则创建
    if not os.path.exists(folder_path):
        os.makedirs(folder_path)
    watermark_image_path = "watermark.png"
    create_watermark_image(watermark_text, watermark_image_path)

    wb_openpyxl = openpyxl.load_workbook(input_file)
    # 创建一个新的 Excel 文件（使用 xlsxwriter）
    wb_xlsxwriter = xlsxwriter.Workbook(output_file)

    # 遍历源文件中的每个工作表
    for sheet_name in wb_openpyxl.sheetnames:
        # 获取当前工作表
        ws_openpyxl = wb_openpyxl[sheet_name]
        # 在新工作簿中创建一个同名的工作表
        ws_xlsxwriter = wb_xlsxwriter.add_worksheet(sheet_name)
        ws_xlsxwriter.set_background(watermark_image_path)

        # 将源工作表中的所有行和列复制到目标工作表
        for row_idx, row in enumerate(ws_openpyxl.iter_rows(), start=0):
            for col_idx, cell in enumerate(row, start=0):
                value = cell.value  # 获取单元格的值
                ws_xlsxwriter.write(row_idx, col_idx, value)

    # 关闭并保存新的工作簿
    wb_xlsxwriter.close()

    # Clean up: remove the temporary watermark image
    os.remove(watermark_image_path)

# Example usage
# add_watermark_to_excel("./data/excel/a.xlsx", "output.xlsx", "四川机密")
